from typing import Any, Generic, Optional, Type, TypeVar, Union, Sequence, Callable, Dict, List
from fastapi_async_sqlalchemy import db
from fastapi.encoders import jsonable_encoder
from sqlmodel import select, func
from pydantic import BaseModel
from sqlalchemy import select, delete, Row, RowMapping, Executable
from sqlalchemy.ext.asyncio import AsyncSession
from sqlmodel.sql.expression import Select
from fastapi_pagination import Params, Page
from fastapi_pagination.ext.async_sqlalchemy import paginate
from models.base_model import SQLModel

T = TypeVar("T", bound=SQLModel)
ModelType = TypeVar('ModelType', bound=SQLModel)
CreateSchemaType = TypeVar('CreateSchemaType', bound=BaseModel)
UpdateSchemaType = TypeVar('UpdateSchemaType', bound=BaseModel)


class CRUDBase(Generic[ModelType]):
    def __init__(self, model: Type[ModelType]):
        self.model = model

    async def get(self, pk: int, db_session: Optional[AsyncSession] = None) -> ModelType | None:
        """
        通过主键 id 获取一条数据
        :param db_session:
        :param pk:
        :return:
        """
        db_session = db_session or db.session
        model = await db_session.execute(select(self.model).where(self.model.id == pk))
        return model.scalars().first()

    async def get_row(
            self, key: str = "id", value: str = None, is_multi=False, db_session: Optional[AsyncSession] = None
    ) -> Sequence[Row[Any] | RowMapping | Any] | ModelType | None | Any:
        """
        获取某条数据
        :param db:
        :param key:
        :param value:
        :param is_multi: 是否返回多条
        :param db_session
        :return:
        """
        db_session = db_session or db.session
        columns = self.model.__table__.columns
        query = select(self.model).where(columns[key] == value)
        response = await db_session.execute(query)
        if is_multi is False:
            return response.scalars().first()
        return response.scalars().all()

    @staticmethod
    async def get_by_query(
            query=None, is_multi=False, db_session: Optional[AsyncSession] = None
    ) -> Sequence[Row[Any] | RowMapping | Any] | ModelType | None | Any:
        db_session = db_session or db.session
        if query is None:
            return None
        response = await db_session.execute(query)
        if is_multi is False:
            return response.scalars().first()
        return response.scalars().all()

    @staticmethod
    async def get_by_query_multi(
            query=None, db_session: Optional[AsyncSession] = None
    ) -> Sequence[Row[Any] | RowMapping | Any] | ModelType | None | Any:
        db_session = db_session or db.session
        if query is None:
            return None
        response = await db_session.execute(query)
        return response.all()

    async def get_count(
            self, query: Optional[Union[T, Select[T]]] = None, db_session: Optional[AsyncSession] = None
    ) -> int:
        db_session = db_session or db.session
        if not query:
            query = select(self.model).subquery()
        response = await db_session.execute(
            select(func.count()).select_from(query)
        )
        return response.scalar_one()

    async def get_multi(
            self,
            skip: int = 0,
            limit: int = 100,
            get_all: bool = True,
            query: Optional[Union[T, Select[T]]] = None,
            db_session: Optional[AsyncSession] = None
    ) -> Sequence[Row[Any] | RowMapping | Any]:
        db_session = db_session or db.session
        if query is None:
            if get_all:
                query = select(self.model).order_by(self.model.id)
            else:
                query = select(self.model).offset(skip).limit(limit).order_by(self.model.id)

        response = await db_session.execute(query)
        return response.scalars().all()

    async def get_multi_paginated(
            self,
            params: Optional[Params] = Params(),
            query: Optional[Union[T, Select[T]]] = None,
            db_session: Optional[AsyncSession] = None
    ) -> Page[ModelType]:
        db_session = db_session or db.session
        if query is None:
            query = select(self.model)
        return await paginate(db_session, query, params)

    async def create(self, obj_in: CreateSchemaType, update: Optional[Dict[str, Any]] = None,
                     db_session: Optional[AsyncSession] = None) -> None:
        """
        新增一条数据

        :param obj_in: Pydantic 模型类
        :param update: Pydantic 模型类
        :param db_session:
        :return:
        """
        db_session = db_session or db.session
        db_obj = self.model.from_orm(obj_in, update=update)  # type: ignore
        db_session.add(db_obj)
        # await db.flush()
        return db_obj

    @staticmethod
    async def update(
            obj_current: ModelType,
            obj_new: UpdateSchemaType | dict[str, Any] | ModelType,
            db_session: Optional[AsyncSession] = None
    ) -> ModelType:
        db_session = db_session or db.session
        obj_data = jsonable_encoder(obj_current)

        if isinstance(obj_new, dict):
            update_data = obj_new
        else:
            update_data = obj_new.dict(
                exclude_unset=True
            )  # This tells Pydantic to not include the values that were not sent
        for field in obj_data:
            if field in update_data:
                setattr(obj_current, field, update_data[field])
        await db_session.merge(obj_current)
        return obj_current

    async def delete(self, pk: int, db_session: Optional[AsyncSession] = None) -> Callable[[], int]:
        """
        通过主键 id 删除一条数据

        :param pk:
        :param db_session:
        :return:
        """
        db_session = db_session or db.session
        result = await db_session.execute(delete(self.model).where(self.model.id == pk))
        return result.rowcount

    def to_sql(self, query: Any = Executable, literal_binds: bool = True):
        """转换为 SQL 语句"""
        return str(query.compile(compile_kwargs={"literal_binds": literal_binds}))

    def to_dicts(self, rows: List[ModelType]) -> List[dict]:
        """Models 转 List[dict]"""
        datas = []
        for row in rows:
            data = {c.name: getattr(row, c.name) for c in row.__table__.columns}
            datas.append(data)
        return datas
